"""
Copyright (c) 2025 Huawei Technologies Co., Ltd.
This file is a part of the CANN Open Software.
Licensed under CANN Open Software License Agreement Version 1.0 (the "License").
Please refer to the License for details. You may not use this file except in compliance with the License.
THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
See LICENSE in the root of the software repository for the full text of the License.
"""
import sys
import numpy as np
import csv
import os

def gen_test_data(times, mode):
    if times < 0:
        print("times must be greater than or equal to 0!")
        sys.exit(1)

    # 保存M N K的路径
    test_dim_csv_filepath = f"./input_params/MNKC_data.csv"

    # 检查并创建 result 目录
    os.makedirs(os.path.dirname(test_dim_csv_filepath), exist_ok=True)

    #生成随机矩阵
    if mode == 0:
        low_zeroPaddingM = 1000
        high_zeroPaddingM = 3000
        low_zeroPaddingN = 1000
        high_zeroPaddingN = 3000
        low_zeroPaddingK = 1000
        high_zeroPaddingK = 3000

        # 表头
        with open(test_dim_csv_filepath, "w") as f_output:
            f_output.write("zeroPaddingM,zeroPaddingN,zeroPaddingK,batchCount\n")

        # 生成数据并写入 CSV 文件
        for i in range(times):

            zeroPaddingM = np.random.randint(low_zeroPaddingM, high_zeroPaddingM, dtype=np.int64)

            zeroPaddingN = np.random.randint(low_zeroPaddingN, high_zeroPaddingN, dtype=np.int64)

            zeroPaddingK = np.random.randint(low_zeroPaddingK, high_zeroPaddingK, dtype=np.int64)

            batchCount = np.random.randint(2, 10, dtype=np.int64)

            with open(test_dim_csv_filepath, "a") as f_output:
                writer = csv.writer(f_output)
                writer.writerow([str(zeroPaddingM), str(zeroPaddingN), str(zeroPaddingK), str(batchCount)])


if __name__ == "__main__":
    times = int(sys.argv[1])  # 测试次数 100
    # mode = int(sys.argv[2])  # 模式 0
    gen_test_data(times, 0)
